// kernel virtual memory management
#include "mem/pmem.h"
#include "mem/vmem.h"
#include "lib/print.h"
#include "lib/str.h"
#include "riscv.h"
#include "memlayout.h"

pgtbl_t kernel_pgtbl; // 内核页表

extern char trampoline[]; // 引用trampoline.S


// 根据pagetable,找到va对应的pte
// 若设置alloc=true 则在PTE无效时尝试申请一个物理页
// 成功返回PTE, 失败返回NULL
// 提示：使用 VA_TO_VPN PTE_TO_PA PA_TO_PTE
// 定义用于获取或创建一个页表项 (PTE) 的函数
pte_t* vm_getpte(pgtbl_t pgtbl, uint64 va, bool alloc) {

    if (va >= VA_MAX)  // 检查虚拟地址是否超出最大允许范围
        panic("vm_getpte: virtual address exceeds maximum!");
    if (!pgtbl)  // 检查页表指针是否为空
        panic("vm_getpte: page table is NULL");

    // 解析虚拟地址的三级页表虚拟页号 (VPN)
    uint64 vpn[3] = {VA_TO_VPN(va, 0), VA_TO_VPN(va, 1), VA_TO_VPN(va, 2)};

    for (int level = 2; level > 0; level--) {
        // 获取当前级别的页表项地址
        pte_t *pte = &pgtbl[vpn[level]];
        // 检查页表项是否有效
        if (*pte & PTE_V) {
            // 如果有效，更新pgtbl指针以指向下一级页表
            pgtbl = (pgtbl_t)PTE_TO_PA(*pte);
        } else {
            // 如果当前页表项无效，根据alloc参数决定是否分配新的物理页
            if (!alloc || (pgtbl = (pgtbl_t)pmem_alloc(false)) == 0) {
                // 如果不分配新页或分配失败，返回NULL
                return NULL;
            }
            // 初始化新分配的物理页（清零）
            memset(pgtbl, 0, PGSIZE);
            // 更新当前页表项为新分配的物理页地址，并设置有效位
            *pte = PA_TO_PTE(pgtbl) | PTE_V;
        }
    }
    // 返回最低级页表项的地址
    return &pgtbl[vpn[0]];
}

// 在pgtbl中建立 [va, va + len) -> [pa, pa + len) 的映射
// 本质是找到va在页表对应位置的pte并修改它
// 检查: va pa 应当是 page-aligned, len(字节数) > 0, va + len <= VA_MAX
// 注意: perm 应该如何使用
void vm_mappages(pgtbl_t pgtbl, uint64 va, uint64 pa, uint64 len, int perm)
{
    // 检查虚拟地址(va)和物理地址(pa)是否都正确对齐到页面大小
    if ((va % PGSIZE) != 0 || (pa % PGSIZE) != 0)
        panic("vm_mappages: VA and PA must be page-aligned");
    // 检查映射长度是否有效
    if (len <= 0)
        panic("vm_mappages: Length must be greater than 0 and a multiple of PGSIZE");
    // 检查虚拟地址范围是否超出最大允许的虚拟地址
    if (va + len > VA_MAX)
        panic("vm_mappages: Virtual address range exceeds the maximum limit");
    
    uint64 last = ALIGN_DOWN(va + len - 1,PGSIZE); // 计算最后一个地址的页面边界
    pte_t *pte;
    // 遍历所需映射的每一页
    for (; va <= last; va += PGSIZE, pa += PGSIZE) {
        pte = vm_getpte(pgtbl, va, true); // 获取或创建页表项
        if (pte == NULL)
            panic("vm_mappages: Failed to get or create PTE");

        // 设置页表项
        *pte = PA_TO_PTE(pa) | perm | PTE_V; // 设置物理地址和权限
    }
}

// 解除pgtbl中[va, va+len)区域的映射
// 如果freeit == true则释放对应物理页, 默认是用户的物理页
void vm_unmappages(pgtbl_t pgtbl, uint64 va, uint64 len, bool freeit) {
    // 检查参数有效性和对齐要求
    if ((va % PGSIZE) != 0 || (len % PGSIZE) != 0 || len <= 0) {
        panic("vm_unmappages: Address and length must be page-aligned and greater than zero");
    }
    if (va + len > VA_MAX) {
        panic("vm_unmappages: Address range exceeds maximum virtual address");
    }

    // 遍历每一页，解除映射并根据情况释放物理页
    for (uint64 a = va; a < va + len; a += PGSIZE) {
        pte_t *pte = vm_getpte(pgtbl, a, false);
        if (!pte || !(*pte & PTE_V)) {
            panic("vm_unmappages: Page not mapped or invalid PTE");
        }
        if (PTE_CHECK(*pte)) {
            panic("vm_unmappages: Attempting to unmap non-leaf PTE");
        }
        if (freeit) {
            uint64 pa = PTE_TO_PA(*pte);
            pmem_free(pa, false);
        }
        *pte = 0;
    }
}

// 完成 UART CLINT PLIC 内核代码区 内核数据区 可分配区域 的映射
// 相当于填充kernel_pgtbl
void kvm_init()
{
    // 为内核的顶级页表分配一个物理页面。
    kernel_pgtbl=(pgtbl_t)pmem_alloc(true);
    memset(kernel_pgtbl, 0, PGSIZE);
    
    // 将内核基地址区域映射到相同的物理地址区域，并提供读、写、执行的权限。
    // 这通常包括内核的文本段（代码）、只读数据等。
    vm_mappages(kernel_pgtbl, (uint64)KERNEL_BASE, (uint64)KERNEL_BASE, (uint64)KERNEL_DATA - (uint64)KERNEL_BASE, PTE_R | PTE_W | PTE_X);

    // 将内核数据段映射到相同的物理地址，只提供读和写权限。
    // 这包括全局变量、静态变量等非执行代码数据。
    vm_mappages(kernel_pgtbl, (uint64)KERNEL_DATA, (uint64)KERNEL_DATA, (uint64)ALLOC_BEGIN - (uint64)KERNEL_DATA, PTE_R | PTE_W);

    // 将动态内存分配区域（堆空间）映射到相同的物理地址，提供读和写权限。
    // 这用于内核运行时的动态内存分配。
    vm_mappages(kernel_pgtbl, (uint64)ALLOC_BEGIN, (uint64)ALLOC_BEGIN, (uint64)ALLOC_END - (uint64)ALLOC_BEGIN, PTE_R | PTE_W);

    // 将 UART（通用异步接收/发送转换器）的基地址映射到相同的物理地址，提供读和写权限。
    // 使内核能够通过编程访问和控制 UART 设备。
    vm_mappages(kernel_pgtbl, UART_BASE, UART_BASE, PGSIZE, PTE_R | PTE_W);

    // 将 CLINT（核间中断控制器）的基地址映射到相同的物理地址，提供读和写权限。
    // 这允许内核配置和响应来自其他处理器的中断。
    vm_mappages(kernel_pgtbl, CLINT_BASE, CLINT_BASE, 0x10000, PTE_R | PTE_W);

    // 将 PLIC（平台级中断控制器）的基地址映射到相同的物理地址，提供读和写权限。
    // 使内核能够管理外部设备（如I/O）发出的中断。
    vm_mappages(kernel_pgtbl, PLIC_BASE, PLIC_BASE, 0x400000, PTE_R | PTE_W);

    // 新增：添加磁盘地址映射
    vm_mappages(kernel_pgtbl,VIRTIO_BASE,VIRTIO_BASE,PGSIZE, PTE_R|PTE_W);

    // 新增：映射 trampoline
    // 将虚拟地址 TRAMPOLINE 映射到物理地址 trampoline，权限设置为可执行（PTE_X）和可读（PTE_R）。
    vm_mappages(kernel_pgtbl, TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_X | PTE_R);

    // 新增：映射 kstack
    // 为第一个内核栈分配一页内存并映射到 KSTACK(1)。
    char *kstack_page = pmem_alloc(true);  // 分配物理内存
    vm_mappages(kernel_pgtbl, KSTACK(1), (uint64)kstack_page, PGSIZE, PTE_R | PTE_W);
}

// 使用新的页表，刷新TLB
void kvm_inithart()
{
    w_satp(MAKE_SATP(kernel_pgtbl));
    sfence_vma();
}

// for debug
// 输出页表内容
void vm_print(pgtbl_t pgtbl)
{
    // 顶级页表，次级页表，低级页表
    pgtbl_t pgtbl_2 = pgtbl, pgtbl_1 = NULL, pgtbl_0 = NULL;
    pte_t pte;

    printf("level-2 pgtbl: pa = %p\n", pgtbl_2);
    for(int i = 0; i < PGSIZE / sizeof(pte_t); i++) 
    {
        pte = pgtbl_2[i];
        if(!((pte) & PTE_V)) continue;
        assert(PTE_CHECK(pte), "vm_print: pte check fail (1)");
        pgtbl_1 = (pgtbl_t)PTE_TO_PA(pte);
        printf(".. level-1 pgtbl %d: pa = %p\n", i, pgtbl_1);
        
        for(int j = 0; j < PGSIZE / sizeof(pte_t); j++)
        {
            pte = pgtbl_1[j];
            if(!((pte) & PTE_V)) continue;
            assert(PTE_CHECK(pte), "vm_print: pte check fail (2)");
            pgtbl_0 = (pgtbl_t)PTE_TO_PA(pte);
            printf(".. .. level-0 pgtbl %d: pa = %p\n", j, pgtbl_2);

            for(int k = 0; k < PGSIZE / sizeof(pte_t); k++) 
            {
                pte = pgtbl_0[k];
                if(!((pte) & PTE_V)) continue;
                assert(!PTE_CHECK(pte), "vm_print: pte check fail (3)");
                printf(".. .. .. physical page %d: pa = %p flags = %d\n", k, (uint64)PTE_TO_PA(pte), (int)PTE_FLAGS(pte));                
            }
        }
    }
}